import os
import argparse
import pandas as pd

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str, default='mllm/images')
    parser.add_argument('--input_file', type=str, default='mllm/images/demo.csv') 
    args = parser.parse_args()

    df = pd.DataFrame(columns=['img_path'])
    df['img_path'] = [os.path.join(args.img_dir, fn) for fn in os.listdir(args.img_dir) if fn.endswith(".png")]

    df.to_csv(args.input_file, index=False) 
    print("csv file saved to: ", args.input_file)